package cn.kickegg.log.core.storage;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.core.IndexRequest;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import cn.kickegg.log.core.config.ElasticsearchConfiguration;
import cn.kickegg.log.core.model.LogEntry;
import cn.kickegg.log.core.spi.LogStorage;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.impl.client.DefaultHttpClient;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.Arrays;
import java.util.UUID;

@Component
@ConditionalOnProperty(name = "log.collector.storage.type", havingValue = "elasticsearch")
public class ElasticsearchLogStorage implements LogStorage {
    
    private final ElasticsearchClient client;
    private final ElasticsearchConfiguration esConfig;
    
    public ElasticsearchLogStorage(ElasticsearchConfiguration esConfig) {
        this.esConfig = esConfig;
        this.client = createClient();
    }
    
    private ElasticsearchClient createClient() {
        RestClientBuilder builder = RestClient.builder(
            Arrays.stream(esConfig.getHosts().split(","))
                .map(host -> {
                    String[] parts = host.split(":");
                    return new HttpHost(parts[0], Integer.parseInt(parts[1]), esConfig.isSsl() ? "https" : "http");
                })
                .toArray(HttpHost[]::new)
        );
        
        // 设置认证信息
        if (esConfig.getUsername() != null && !esConfig.getUsername().isEmpty()) {
            BasicCredentialsProvider credentialsProvider = new BasicCredentialsProvider();
            credentialsProvider.setCredentials(
                AuthScope.ANY,
                new UsernamePasswordCredentials(esConfig.getUsername(), esConfig.getPassword())
            );
            builder.setHttpClientConfigCallback(httpClientBuilder ->
                httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider)
            );
        }
        
        // 设置超时
        builder.setRequestConfigCallback(requestConfigBuilder ->
            requestConfigBuilder
                .setConnectTimeout(esConfig.getConnectionTimeout())
                .setSocketTimeout(esConfig.getSocketTimeout())
        );
        
        RestClient restClient = builder.build();
        RestClientTransport transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
        return new ElasticsearchClient(transport);
    }
    
    @Override
    public void store(LogEntry logEntry) {
        try {
            logEntry.setId(UUID.randomUUID().toString());
            IndexRequest<LogEntry> request = IndexRequest.of(i -> 
                i.index(esConfig.getIndexName())
                 .id(logEntry.getId())
                 .document(logEntry)
            );
            
            client.index(request);
        } catch (IOException e) {
            throw new RuntimeException("Failed to store log in Elasticsearch", e);
        }
    }
    
    @Override
    public String getType() {
        return "elasticsearch";
    }
} 